from flask import Flask, request, jsonify, render_template, session
import mysql.connector
from mysql.connector import Error
from pathlib import Path
import docx
import numpy as np
import pandas as pd
import tempfile
import fitz  # PyMuPDF
import jieba
import jieba.posseg as pseg  # 用于词性标注
from aip import AipOcr
import pdfplumber
from io import BytesIO
import os
import json
import requests
import torch
from transformers import BertTokenizer, BertModel

app = Flask("MyApp")
app.secret_key = os.urandom(24)  # 用于加密会话数据

# 数据库配置
db_config = {
    "host": "localhost",  # 数据库地址
    "user": "root",  # 数据库用户名
    "password": "",  # 数据库密码
    "database": "tianshi",  # 数据库名
    "charset": "utf8"
}

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# 加载 BERT 模型和 Tokenizer
EMBEDDING_MODEL = 'E:/model6/zpoint_large_embedding_zh/'  # 替换为实际路径
tokenizer = BertTokenizer.from_pretrained(EMBEDDING_MODEL)
bert_model = BertModel.from_pretrained(EMBEDDING_MODEL)
bert_model.to(device)

# 停用词列表
STOP_WORDS = {"怎么", "了", "是", "什么", "的", "如何", "吗", "呢", "啊", "吧", "呀"}

# 连接数据库
def get_db_connection():
    try:
        connection = mysql.connector.connect(**db_config)
        if connection.is_connected():
            return connection
    except Error as e:
        print("数据库连接失败:", e)
        return None

# 计算嵌入或加载现有嵌入
def calculate_embeddings(data, output_json_file):
    if os.path.exists(output_json_file):
        print(f"发现现有的嵌入文件：{output_json_file}，跳过计算。")
        with open(output_json_file, 'r', encoding='utf-8') as json_file:
            return json.load(json_file)

    # 初始化结果列表
    articles = []

    # 设置模型为评估模式
    bert_model.eval()
    with torch.no_grad():
        for row in data:
            title = row['title']
            link = row['link']

            # 计算 BERT 嵌入
            inputs = tokenizer(title, return_tensors='pt', padding=True, truncation=True, max_length=512)
            inputs = {key: value.to(device) for key, value in inputs.items()}  # 移动到 GPU 或 CPU
            outputs = bert_model(**inputs)
            cls_embedding = outputs.last_hidden_state[:, 0, :].squeeze().cpu().numpy()  # 获取 [CLS] 嵌入

            # 存储结果
            articles.append({
                "title": title,
                "link": link,
                "embedding": cls_embedding.tolist()
            })

    # 保存为 JSON 文件
    with open(output_json_file, 'w', encoding='utf-8') as json_file:
        json.dump(articles, json_file, ensure_ascii=False, indent=4)

    print(f"所有嵌入已保存到 {output_json_file}")
    return articles

# 检索相似文章
def retrieve_similar_articles(query, articles, top_n=5):
    bert_model.eval()
    with torch.no_grad():
        inputs = tokenizer(query, return_tensors='pt', padding=True, truncation=True, max_length=512)
        inputs = {key: value.to(device) for key, value in inputs.items()}
        outputs = bert_model(**inputs)
        query_embedding = outputs.last_hidden_state[:, 0, :].squeeze().cpu().numpy()

    def cosine_similarity(vec1, vec2):
        dot_product = np.dot(vec1, vec2)
        norm_vec1 = np.linalg.norm(vec1)
        norm_vec2 = np.linalg.norm(vec2)
        return dot_product / (norm_vec1 * norm_vec2)

    similarities = []
    for article in articles:
        similarity = cosine_similarity(query_embedding, np.array(article['embedding']))
        similarities.append((article, similarity))

    similarities = sorted(similarities, key=lambda x: x[1], reverse=True)

    seen = set()
    results = []
    for article, similarity in similarities:
        if article['title'] not in seen:
            seen.add(article['title'])
            results.append({
                "title": article['title'],
                "url": article['link'],
                "score": similarity,
            })
        if len(results) == top_n:
            break

    return results

# 推荐路由
@app.route('/recommendations', methods=['GET'])
def recommendations():
    query = request.args.get('query', '')

    if not query:
        return jsonify({"message": "输入不能为空", "articles": []}), 400

    connection = get_db_connection()
    if not connection:
        return jsonify({"message": "无法连接到数据库"}), 500

    try:
        cursor = connection.cursor(dictionary=True)
        cursor.execute("SELECT title, link FROM data")
        rows = cursor.fetchall()

        # 计算嵌入或加载现有嵌入
        output_json_file = 'article_embeddings.json'
        articles = calculate_embeddings(rows, output_json_file)

        # 检索相似文章
        recommendations = retrieve_similar_articles(query, articles, top_n=5)

        if recommendations:
            return jsonify({"message": "推荐成功", "articles": recommendations}), 200
        else:
            return jsonify({"message": "没有找到相关的文章", "articles": []}), 404

    except Exception as e:
        return jsonify({"message": f"数据库查询失败: {str(e)}"}), 500

# 文件处理功能
def extract_text_from_docx(file_path):
    doc = docx.Document(file_path)
    return "\n".join([para.text for para in doc.paragraphs])

def extract_text_from_excel(file_path):
    df = pd.read_excel(file_path)
    return df.to_string(index=False)

def extract_text_from_pdf(file_path):
    text = ""
    with fitz.open(file_path) as pdf_doc:
        for page_num in range(pdf_doc.page_count):
            page = pdf_doc[page_num]
            text += page.get_text("text")
    return text


def get_image_content(file_path):
    with open(file_path, 'rb') as fp:
        return fp.read()


def get_pdf_pages_as_images(pdf_path, dpi=300):
    images = []
    with pdfplumber.open(pdf_path) as pdf:
        for page in pdf.pages:
            pil_image = page.to_image(resolution=dpi).original
            buffered = BytesIO()
            pil_image.save(buffered, format="PNG")
            images.append(buffered.getvalue())
    return images


def perform_ocr(image_data):
    options = {
        "language_type": "CHN_ENG",
        "detect_direction": "true",
        "detect_language": "true",
        "probability": "true"
    }
    result = client.basicGeneral(image_data, options)
    words_only = [item['words'] for item in result.get('words_result', [])]
    return "\n".join(words_only)


def ocr_file(file_path):
    global ocr_text_content
    try:
        if file_path.lower().endswith(('.png', '.jpg', '.jpeg')):
            image_data = get_image_content(file_path)
            ocr_text_content = perform_ocr(image_data)
        elif file_path.lower().endswith('.pdf'):
            images = get_pdf_pages_as_images(file_path, dpi=300)
            ocr_text_content = ""
            for idx, image in enumerate(images):
                page_text = perform_ocr(image)
                ocr_text_content += f"Page {idx + 1}:\n{page_text}\n\n"
        else:
            raise ValueError("Unsupported file type.")
    except Exception as e:
        raise ValueError(f"File processing failed: {e}")


# 路由: 提供登录页面
@app.route('/')
def login_page():
    return render_template('login.html')  # 返回 login.html 页面

# 路由: 登录用户
@app.route('/login', methods=['POST'])
def login_user():
    data = request.json
    email = data.get('email')
    password = data.get('password')

    if not email or not password:
        return jsonify({"message": "邮箱和密码不能为空"}), 400

    connection = get_db_connection()
    if not connection:
        return jsonify({"message": "无法连接到数据库"}), 500

    try:
        cursor = connection.cursor()

        cursor.execute("SELECT * FROM user WHERE email = %s", (email,))
        user = cursor.fetchone()

        if not user:
            return jsonify({"message": "用户不存在"}), 404

        stored_password = user[2]  # 假设密码存储在第三列
        if password != stored_password:
            return jsonify({"message": "密码错误"}), 401

        session['email'] = email  # 登录成功后将邮箱存入会话
        return jsonify({"message": "登录成功", "redirect": "/index"}), 200

    except Error as e:
        print("数据库操作失败:", e)
        return jsonify({"message": "登录失败"}), 500

    finally:
        if connection.is_connected():
            cursor.close()
            connection.close()

@app.route('/register')
def index():
    return render_template('register.html')  # 返回 register.html 页面

# 路由: 注册用户
@app.route('/register', methods=['POST'])
def register_user():
    data = request.json
    email = data.get('email')
    password = data.get('password')

    if not email or not password:
        return jsonify({"message": "邮箱和密码不能为空"}), 400

    connection = get_db_connection()
    if not connection:
        return jsonify({"message": "无法连接到数据库"}), 500

    try:
        cursor = connection.cursor()

        cursor.execute("SELECT * FROM user WHERE email = %s", (email,))
        existing_user = cursor.fetchone()

        if existing_user:
            return jsonify({"message": "该邮箱已注册"}), 409

        insert_query = "INSERT INTO user (email, password) VALUES (%s, %s)"
        cursor.execute(insert_query, (email, password))
        connection.commit()

        return jsonify({"message": "注册成功", "redirect": "/login"}), 201

    except Error as e:
        print("数据库操作失败:", e)
        return jsonify({"message": "注册失败"}), 500

    finally:
        if connection.is_connected():
            cursor.close()
            connection.close()

# 路由: 提供忘记密码页面
@app.route('/change')
def show_reset_password_page():
    return render_template('change.html')  # 返回 change.html 页面

# 路由: 重置密码
@app.route('/change', methods=['POST'])
def reset_password():
    data = request.json
    email = data.get('email')
    new_password = data.get('newPassword')

    # 前端输入校验
    if not email or not new_password:
        return jsonify({"message": "邮箱和新密码不能为空"}), 400

    # 连接数据库
    connection = get_db_connection()
    if not connection:
        return jsonify({"message": "无法连接到数据库"}), 500

    try:
        cursor = connection.cursor()

        # 检查用户是否存在
        cursor.execute("SELECT * FROM user WHERE email = %s", (email,))
        user = cursor.fetchone()

        if not user:
            return jsonify({"message": "该邮箱未注册"}), 404

        # 更新密码
        update_query = "UPDATE user SET password = %s WHERE email = %s"
        cursor.execute(update_query, (new_password, email))
        connection.commit()

        return jsonify({"message": "密码重置成功"}), 200

    except Error as e:
        print("数据库操作失败:", e)
        return jsonify({"message": "密码重置失败"}), 500

    finally:
        if connection.is_connected():
            cursor.close()
            connection.close()
# 路由: 获取用户的历史对话
@app.route('/get-history', methods=['GET'])
def get_history():
    if 'email' not in session:
        return jsonify({"message": "用户未登录"}), 403

    user_email = session['email']
    connection = get_db_connection()

    if not connection:
        return jsonify({"message": "无法连接到数据库"}), 500

    try:
        cursor = connection.cursor(dictionary=True)  # 使用字典形式获取结果
        cursor.execute(
            "SELECT id, question, answer, timestamp FROM user_conversations WHERE email = %s ORDER BY timestamp DESC",
            (user_email,)
        )
        history = cursor.fetchall()
        return jsonify({"message": "获取历史记录成功", "history": history}), 200

    except Error as e:
        print("获取历史记录失败:", e)
        return jsonify({"message": "获取历史记录失败"}), 500

    finally:
        if connection.is_connected():
            cursor.close()
            connection.close()

# 路由: 提供主页
@app.route('/index')
def index_page():
    return render_template('index.html')  # 返回 index.html 页面
@app.route("/chatgpt-clone", methods=["POST"])
def chatgpt_clone():
    global ocr_text_content
    if 'email' not in session:
        return jsonify({"message": "请先登录"}), 403  # 检查用户是否已登录

    user_email = session['email']
    question = request.form.get("question", "").strip()
    file = request.files.get("file")

    # 初始化会话上下文
    if 'messages' not in session:
        session['messages'] = [{"role": "assistant", "content": "我是中医问答助手天士，很高兴为您解答中医相关问题。"}]

    # 处理上传文件
    if file:
        file_extension = Path(file.filename).suffix.lower()
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp:
                file.save(tmp.name)
                if file_extension == ".docx":
                    ocr_text_content = extract_text_from_docx(tmp.name)
                elif file_extension == ".xlsx":
                    ocr_text_content = extract_text_from_excel(tmp.name)
                elif file_extension == ".pdf":
                    ocr_text_content = extract_text_from_pdf(tmp.name)  # Extract text directly from PDF
                    if not ocr_text_content:  # 如果直接提取的文本为空，尝试OCR
                        ocr_file(tmp.name)
                elif file_extension in [".png", ".jpg", ".jpeg"]:
                    ocr_file(tmp.name)
                else:
                    return "Unsupported file type. Please upload a .docx, .xlsx, .pdf, .png, .jpg, or .jpeg file."

            # 将 OCR 内容和用户问题拼接，并加到上下文中
            combined_question = f"请根据以下内容回答问题：\n{ocr_text_content}\n{question}"
            session['messages'].append({"role": "user", "content": combined_question})

        except Exception as e:
            return f"Error processing file: {str(e)}"
    else:
        # 如果没有上传文件，仅处理用户输入的问题
        session['messages'].append({"role": "user", "content": question})

    if question:
        try:
            # 检查用户是否询问助手身份
            if "你是谁" in question:
                answer = "我是中医问答助手天士，专为您解答中医相关问题。"
                session['messages'].append({"role": "assistant", "content": answer})
                return answer

            # 发送请求到模型 API (需要将 URL 替换为实际模型服务的地址)
            model_api_url = "http://localhost:11434/api/chat"  # 替换为新模型服务 URL
            response = requests.post(model_api_url, headers={'Content-Type': 'application/json'}, data=json.dumps({
                "model": "qwen2.5:latest",
                "options": {
                    "temperature": 0.0  # 设置输出更加确定
                },
                "stream": False,
                "messages": session['messages']
            }))

            # 解析模型的响应
            response_json = response.json()
            response_content = response_json.get('message', {}).get('content', 'No response from model')
            session['messages'].append({"role": "assistant", "content": response_content})

            # 存储问答记录到数据库
            try:
                connection = get_db_connection()
                if connection:
                    cursor = connection.cursor()
                    insert_query = """
                    INSERT INTO user_conversations (email, question, answer) 
                    VALUES (%s, %s, %s)
                    """
                    cursor.execute(insert_query, (user_email, question, response_content))
                    connection.commit()
                    cursor.close()
                    connection.close()
            except Exception as db_error:
                return f"数据库错误: {str(db_error)}"

            return response_content
        except Exception as e:
            return f"错误: {str(e)}"

    return "没有输入内容"


if __name__ == "__main__":
    app.run(debug=True)
